import modelscope
import datasets
import transformers
import os
import torch
from torch.utils.data import Dataset
from torch.utils.data import random_split
from torch.utils.data import DataLoader
# from torch.utils.data.distributed import DistributedSampler
# from torch.nn.parallel import DistributedDataParallel as DDP
# import torch.distributed as dist
# dist.init_process_group(backend="nccl")

# Accelerate框架封装后的DDP执行分布式训练
from accelerate import Accelerator

## 原生的DDP执行并行训练的配置很复杂，一旦出现配置错误也很难debug，基于accelerate框架要稳定很多
## execute cmd
# cd root_path
# python -m torch.distributed.launch --nproc_per_node=8 ./scripts/pretrain/pretrain_with_pytorch_qwen2_5_0_5B.py

# Accelerate框架封装后的DDP执行分布式训练
## execute cmd
# cd root_path
# accelerate launch ./scripts/pretrain/pretrain_with_pytorch_qwen2_5_0_5B.py 

## path define
root_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
filename = os.path.basename(__file__).split('.')[0]
dataset_path = os.path.join(root_path, "datasets/pretrain/wiki-datasets/wikipedia-zh-cn-20241020.json")
model_ID = 'Qwen/Qwen2.5-0.5B'

model_path = os.path.join(root_path, 'llms/' + model_ID)
checkpoint_path = os.path.join(root_path,"checkpoints/"+filename)
model_save_path = os.path.join(root_path,'llms/' + filename)
log_path = os.path.join(root_path,'logs/' + filename)

# initial model and tokenizer
modelscope.AutoConfig.from_pretrained(model_ID).save_pretrained(model_path)
modelscope.AutoTokenizer.from_pretrained(model_ID).save_pretrained(model_path)
tokenizer = transformers.AutoTokenizer.from_pretrained(model_path)
tokenizer.pad_token = tokenizer.eos_token
context_length = 64

# Accelerate框架封装后的DDP执行分布式训练
accelerator = Accelerator()

class MyDataset(Dataset):
    def __init__(self) -> None:
        super().__init__()
        self.data = datasets.load_dataset("json", data_files = dataset_path)['train']


    # return: text | label
    def __getitem__(self, index):
        return self.data[index]['text'], self.data[index]["title"]
    
    def __len__(self):
        return len(self.data)

def prepare_dataloader():
    dataset = MyDataset()
    trainset, validset = random_split(
        dataset=dataset, 
        lengths=[0.9, 0.1], 
        generator=torch.Generator().manual_seed(3000))
    
    def tokenize(batch):
        texts, labels = [], []
        for item in batch:
            texts.append(item[0])
            labels.append(item[1])
        input_tokens = tokenizer(texts, max_length=context_length, padding=True, truncation=True, return_tensors="pt")
        labels_tokens = tokenizer(labels, padding='longest', truncation=True, return_tensors="pt")
        input_ids = torch.cat((input_tokens['input_ids'], labels_tokens['input_ids'], torch.ones(len(texts), 1)*int(tokenizer.pad_token_id)), dim=1)
        attention_mask = torch.cat((input_tokens['attention_mask'], labels_tokens['attention_mask'], torch.ones(len(texts), 1)), dim=1)
        labels = torch.cat((-100 * torch.ones(input_tokens['input_ids'].shape), labels_tokens['input_ids'], torch.ones(len(texts), 1)*int(tokenizer.pad_token_id)), dim=1)
    
        return {'input_ids': torch.LongTensor(input_ids.numpy()), 
            'attention_mask': torch.LongTensor(attention_mask.numpy()),
            'labels': torch.LongTensor(labels.numpy())}
    
    # 原生的pytorch DDP执行分布式训练
    # trainloader = DataLoader(trainset, batch_size=32, collate_fn=tokenize, sampler=DistributedSampler(trainset))
    # validloader = DataLoader(validset, batch_size=64, collate_fn=tokenize, sampler=DistributedSampler(validset))

    # Accelerate框架封装后的DDP执行分布式训练
    trainloader = DataLoader(trainset, batch_size=32, collate_fn=tokenize, shuffle=True)
    validloader = DataLoader(validset, batch_size=64, collate_fn=tokenize, shuffle=True)
    return trainloader, validloader

def prepare_model_and_optimizer():
    config = transformers.AutoConfig.from_pretrained(
        model_path,
        vocab_size=len(tokenizer),
        hidden_size=256,
        intermediate_size=2048,
        num_attention_heads=8,
        num_hidden_layers=12,
        n_ctx=context_length,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,)
    print("Model Config:")
    print(config)

    model = transformers.Qwen2ForCausalLM(config)
   
    # 原生的pytorch DDP执行分布式训练
    # send model to device here
    # if torch.cuda.is_available():
    #     model = model.to(int(os.environ["LOCAL_RANK"]))
    
    # model = DDP(model)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-5)
    return model, optimizer

def evaluate(model, validloader):
    model.eval()
    acc_num = 0
    with torch.inference_mode():
        for batch in validloader:
            # 原生的pytorch DDP执行分布式训练
            # if torch.cuda.is_available():
            #     batch = {k: v.to(int(os.environ["LOCAL_RANK"])) for k, v in batch.items()}
            output = model(**batch)
            pred = torch.argmax(output.logits, dim=-1)

            # Accelerate框架封装后的DDP执行分布式训练
            # 将所有机器的预测结果进行汇总
            # gather_for_metrics是accelerator的内置方法，用于汇总通信组的信息
            pred, refs = accelerator.gather_for_metrics((pred, batch["labels"]))
            acc_num += (pred.long() == refs.long()).float().sum()

             # 原生的pytorch DDP执行分布式训练
            #acc_num += (pred.long() == batch["labels"].long()).float().sum()
    # dist.all_reduce(acc_num)
    return acc_num / len(validloader.dataset)

# 原生的pytorch DDP执行分布式训练
# def print_rank_0(info):
#     if int(os.environ["RANK"]) == 0:
#         print(info)

def train(model, optimizer, trainloader, validloader, epoch=3, log_step=100):
    global_step = 0
    for ep in range(epoch):
        model.train()
        # 原生的pytorch DDP执行分布式训练
        # trainloader.sampler.set_epoch(ep)
        for batch in trainloader:
            # 原生的pytorch DDP执行分布式训练
            # if torch.cuda.is_available():
            #     # send tensor to device here
            #     batch = {k: v.to(int(os.environ["LOCAL_RANK"])) for k, v in batch.items()}
            optimizer.zero_grad()
            output = model(**batch)
            loss = output.loss
            # Accelerate框架封装后的DDP执行分布式训练
            accelerator.backward(loss)
            # 原生的pytorch DDP执行分布式训练
            # loss.backward()
            optimizer.step()
            if global_step % log_step == 0:
                # 原生的pytorch DDP执行分布式训练
                # dist.all_reduce(loss, op=dist.ReduceOp.AVG)
                # print_rank_0(f"ep: {ep}, global_step: {global_step}, loss: {loss.item()}")

                # 将Loss在所有机器上合并取均值，不然不同机器的Loss是不一样的
                # accelerator提供了同样的包装
                loss = accelerator.reduce(loss, "mean")
                # 直接可以print日志，而不需要指定Rank来print
                accelerator.print(f"ep: {ep}, global_step: {global_step}, loss: {loss.item()}")
            global_step += 1
        acc = evaluate(model, validloader)
        accelerator.print(f"ep: {ep}, acc: {acc}")
        # 原生的pytorch DDP执行分布式训练
        # print_rank_0(f"ep: {ep}, acc: {acc}")   

def train_qwen_with_pytorch_on_distributed_env():
    trainloader, validloader = prepare_dataloader()
    model, optimizer = prepare_model_and_optimizer()

     # Accelerate框架封装后的DDP执行分布式训练
    model, optimizer, trainloader, validloader = accelerator.prepare(model, optimizer, trainloader, validloader)

    train(model, optimizer, trainloader, validloader)

if __name__ == "__main__":
    train_qwen_with_pytorch_on_distributed_env()






